package main

import (
	"crypto/rand"
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"os"
	"time"
)

// PerformanceResult (已更新)
type PerformanceResult struct {
	Experiment     string
	Parameter      string
	Value          int
	ProofGenTimeMs float64
	VerifyTimeMs   float64
	// SetupTimeMs         float64
	// AuditReqTimeMs      float64
	// AvgProofGenTimeMs   float64
	// TotalProofGenTimeMs float64
	// AvgVerifyTimeMs     float64
	// TotalVerifyTimeMs   float64
}

// generateDummyFile 在内存中创建随机数据
func generateDummyFile(size int) ([]byte, error) {
	data := make([]byte, size)
	if _, err := rand.Read(data); err != nil {
		return nil, err
	}
	return data, nil
}

// 基本功能性实验
func basicExp() {
	log.Println("--- Starting Decentralized Multi-Cloud Auditing Simulation ---")

	sys := NewSystem(2 * 1024)
	dataID := "my_secret_document.txt"
	log.Println("--- System Setup Complete ---")

	dataOwner := NewDataOwner("Alan")
	cspA := NewCSP("CSPA")
	cspB := NewCSP("CSPB")
	cspC_malicious := NewCSP("CSP_Malicious")
	allCSPs := []*CSP{cspA, cspB, cspC_malicious}

	blockchain := NewBlockChain(sys)
	for _, c := range allCSPs {
		c.GenerateKeyPair(sys)
		blockchain.RegisterCSP(c)
	}
	log.Println("--- Entities creation and registration Complete ---")

	originalData, _ := generateDummyFile(2 * 1024 * 1024)
	corruptedData := make([]byte, len(originalData))
	copy(corruptedData, originalData)
	corruptedData[100]++

	dataOwner.UploadData(dataID, originalData, []*CSP{cspA, cspB})
	dataOwner.UploadData(dataID, corruptedData, []*CSP{cspC_malicious})
	log.Println("--- Data Upload to CSPs' memory Complete ---")

	tau := dataOwner.RequestVerification(dataID)
	log.Printf("--- Auditing Request Sent (τ: %s) ---", tau)

	proofs := make(map[string]*Proof)
	for _, c := range allCSPs {
		proof, err := c.GenerateProof(dataID, tau, sys, allCSPs)
		if err != nil {
			log.Fatalf("CSP %s failed to generate proof: %v", c.Name, err)
		}
		proofs[c.Name] = proof
	}
	log.Println("--- All CSPs generated proofs ---")

	fmt.Println("\n--- SCENARIO 1: Verifying two honest CSPs (A and B) ---")
	resultAB := blockchain.Verify(proofs[cspA.Name], proofs[cspB.Name], cspA.Name, cspB.Name)
	if resultAB {
		fmt.Printf("✅ VERIFICATION SUCCESS: Data in %s and %s is consistent.\n", cspA.Name, cspB.Name)
	} else {
		fmt.Printf("❌ VERIFICATION FAILED: Data in %s and %s is INCONSISTENT.\n", cspA.Name, cspB.Name)
	}

	fmt.Println("\n--- SCENARIO 2: Verifying an honest and a malicious CSP (A and C) ---")
	resultAC := blockchain.Verify(proofs[cspA.Name], proofs[cspC_malicious.Name], cspA.Name, cspC_malicious.Name)
	if !resultAC {
		fmt.Printf("✅ INCONSISTENCY DETECTED: As expected, data between %s and %s is inconsistent.\n", cspA.Name, cspC_malicious.Name)
	} else {
		fmt.Printf("❌ DETECTION FAILED: Inconsistency was NOT detected between %s and %s.\n", cspA.Name, cspC_malicious.Name)
	}
}

// 实验1：考察文件大小的影响
func fileSizeExp(writer *csv.Writer) {
	log.Println("\n\n<<<<<<<<<< Demo Experiment 1: Performance vs. File Size >>>>>>>>>>")
	const K = 5
	const NumBlocks = 4096
	fileSizeOptionsMB := []int{4, 8, 12, 16, 20}

	for _, sizeMB := range fileSizeOptionsMB {
		fileSize := sizeMB * 1024 * 1024
		blockSize := fileSize / NumBlocks
		log.Printf("\n--- Testing with FileSize: %dMB ---", sizeMB)

		sys := NewSystem(blockSize)
		dataOwner := NewDataOwner("Alan")
		allCSPs := make([]*CSP, K)
		for i := 0; i < K; i++ {
			allCSPs[i] = NewCSP(fmt.Sprintf("CSP-%d", i))
		}
		bc := NewBlockChain(sys)

		for _, c := range allCSPs {
			c.GenerateKeyPair(sys)
			bc.RegisterCSP(c)
		}

		dummyFile, _ := generateDummyFile(fileSize)
		dataOwner.UploadData("perf_doc", dummyFile, allCSPs)

		tau := dataOwner.RequestVerification("perf_doc")

		// TODO：生成k-1个proof，但是我是定义在一个函数里的
		start := time.Now()
		proof, _ := allCSPs[0].GenerateProof("perf_doc", tau, sys, allCSPs)
		proofGenTime := time.Since(start)

		// TODO: 新方案也是两两验证，所以是k-1次
		start = time.Now()
		for range K - 1 {
			bc.Verify(proof, proof, allCSPs[0].Name, allCSPs[1].Name)
		}
		verifyTime := time.Since(start)

		result := PerformanceResult{
			Experiment: "FileSize", Parameter: "FileSize(MB)", Value: sizeMB,
			ProofGenTimeMs: float64(proofGenTime.Microseconds()) / 1000.0,
			VerifyTimeMs:   float64(verifyTime.Microseconds()) / 1000.0,
			// SetupTimeMs:         float64(avgSetupTime.Microseconds()) / 1000.0,
			// AuditReqTimeMs:      float64(auditReqTime.Microseconds()) / 1000.0,
			// AvgProofGenTimeMs:   float64(avgProofGenTime.Microseconds()) / 1000.0,
			// TotalProofGenTimeMs: float64(totalProofGenTime.Microseconds()) / 1000.0,
			// AvgVerifyTimeMs:     float64(verifyTime.Microseconds()) / 1000.0,
			// TotalVerifyTimeMs:   float64(verifyTime.Microseconds()) / 1000.0,
		}
		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs),
			fmt.Sprintf("%.3f", result.VerifyTimeMs),
			// fmt.Sprintf("%.3f", result.SetupTimeMs), fmt.Sprintf("%.3f", result.AuditReqTimeMs),
			// fmt.Sprintf("%.3f", result.AvgProofGenTimeMs), fmt.Sprintf("%.3f", result.TotalProofGenTimeMs),
			// fmt.Sprintf("%.3f", result.AvgVerifyTimeMs), fmt.Sprintf("%.3f", result.TotalVerifyTimeMs),
		})
		writer.Flush()
	}
}

// 实验2：考察 CSP 数量的影响
func cspCountExp(writer *csv.Writer) {
	log.Println("\n\n<<<<<<<<<< Demo Experiment 2: Performance vs. Number of CSPs >>>>>>>>>>")
	const NumBlocks = 4096
	const BlockSize = 1024
	const FileSize = NumBlocks * BlockSize
	cspCounts := []int{5, 10, 15, 20, 25, 30}

	dummyFile, _ := generateDummyFile(FileSize)

	for _, k := range cspCounts {
		log.Printf("\n--- Testing with %d CSPs ---", k)

		sys := NewSystem(BlockSize)
		dataOwner := NewDataOwner("Alan")
		allCSPs := make([]*CSP, k)
		for i := 0; i < k; i++ {
			allCSPs[i] = NewCSP(fmt.Sprintf("CSP-%d", i))
		}
		bc := NewBlockChain(sys)

		for _, c := range allCSPs {
			c.GenerateKeyPair(sys)
			bc.RegisterCSP(c)
		}

		dataOwner.UploadData("perf_doc", dummyFile, allCSPs)

		tau := dataOwner.RequestVerification("perf_doc")

		// TODO：生成k-1个proof，但是我是定义在一个函数里的
		start := time.Now()
		proof, _ := allCSPs[0].GenerateProof("perf_doc", tau, sys, allCSPs)
		proofGenTime := time.Since(start)

		// TODO: 新方案也是两两验证，所以是k-1次
		start = time.Now()
		for range k - 1 {
			bc.Verify(proof, proof, allCSPs[0].Name, allCSPs[1].Name)
		}
		verifyTime := time.Since(start)
		result := PerformanceResult{
			Experiment: "CSPCount", Parameter: "NumCSPs", Value: k,
			ProofGenTimeMs: float64(proofGenTime.Microseconds()) / 1000.0,
			VerifyTimeMs:   float64(verifyTime.Microseconds()) / 1000.0,
		}
		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs),
			fmt.Sprintf("%.3f", result.VerifyTimeMs),
		})
		writer.Flush()
	}
}

// 实验3：考察文件分块数量的影响
func numBlocksExp(writer *csv.Writer) {
	log.Println("\n\n<<<<<<<<<< Demo Experiment 3: Performance vs. Number of Blocks >>>>>>>>>>")
	const K = 5
	const FileSize = 4 * 1024 * 1024
	numBlocksOptions := []int{1024, 2048, 4096, 8192, 16384}

	dummyFile, _ := generateDummyFile(FileSize)

	for _, numBlocks := range numBlocksOptions {
		blockSize := FileSize / numBlocks
		log.Printf("\n--- Testing with %d blocks ---", numBlocks)

		sys := NewSystem(blockSize)
		dataOwner := NewDataOwner("Alan")
		allCSPs := make([]*CSP, K)
		for i := 0; i < K; i++ {
			allCSPs[i] = NewCSP(fmt.Sprintf("CSP-%d", i))
		}
		bc := NewBlockChain(sys)

		for _, c := range allCSPs {
			c.GenerateKeyPair(sys)
			bc.RegisterCSP(c)
		}

		dataOwner.UploadData("perf_doc", dummyFile, allCSPs)

		tau := dataOwner.RequestVerification("perf_doc")

		// TODO：生成k-1个proof，但是我是定义在一个函数里的
		start := time.Now()
		proof, _ := allCSPs[0].GenerateProof("perf_doc", tau, sys, allCSPs)
		proofGenTime := time.Since(start)

		// TODO: 新方案也是两两验证，所以是k-1次
		start = time.Now()
		for range K - 1 {
			bc.Verify(proof, proof, allCSPs[0].Name, allCSPs[1].Name)
		}
		verifyTime := time.Since(start)
		
		result := PerformanceResult{
			Experiment: "NumBlocks", Parameter: "NumberOfBlocks", Value: numBlocks,
			ProofGenTimeMs: float64(proofGenTime.Microseconds()) / 1000.0,
			VerifyTimeMs:   float64(verifyTime.Microseconds()) / 1000.0,
		}
		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs),
			fmt.Sprintf("%.3f", result.VerifyTimeMs),
		})
		writer.Flush()
	}
}

func main() {
	logFile, err := os.OpenFile("experiment_demo.log", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
	if err != nil {
		log.Fatalf("Failed to open log file: %v", err)
	}
	defer logFile.Close()
	mw := io.MultiWriter(os.Stdout, logFile)
	log.SetOutput(mw)

	csvFile, err := os.Create("results_demo.csv")
	if err != nil {
		log.Fatalf("Failed to create CSV file: %v", err)
	}
	defer csvFile.Close()
	csvWriter := csv.NewWriter(csvFile)
	defer csvWriter.Flush()
	csvWriter.Write([]string{"Experiment", "Parameter", "Value", "ProofGenTime(ms)", "VerifyTime(ms)"})

	// basicExp()
	fileSizeExp(csvWriter)
	cspCountExp(csvWriter)
	numBlocksExp(csvWriter)

	log.Println("\n\n<<<<<<<<<< All experiments for demo finished. >>>>>>>>>>")
}
